from typing import List
import os
from pathlib import Path

from setuptools import Extension, setup, find_packages
from setuptools.command.build_ext import build_ext

TARGET_DEVICE = "npu"
ROOT_DIR = os.path.dirname(__file__)

def _is_npu() -> bool:
    return TARGET_DEVICE == "npu"

class CMakeExtension(Extension):
    def __init__(self, name: str, src_dir: str = "") -> None:
        super().__init__(name, sources=[])
        self.src_dir = os.fspath(Path(src_dir).resolve())

class CMakeBuild(build_ext):
    user_options = build_ext.user_options[:]
    user_options += [
        ('cmake-args=', None, 'Arguments to pass to CMake'),
    ]

    def initialize_options(self):
        super().initialize_options()
        self.cmake_args = None

    def finalize_options(self):
        super().finalize_options()
        if self.cmake_args is None:
            self.cmake_args = ""
    
    # def command_consumes_arguments(self):
    #     return True

    def run(self):
        for ext in self.extensions:
            if isinstance(ext, CMakeExtension):
                self.cmake_build(ext)
        super().run()

    def cmake_build(self, ext):
        cwd = Path().absolute()
        build_temp = f"{Path(self.build_temp)}/{ext.name}"
        os.makedirs(build_temp, exist_ok=True)
        ext_dir = Path(self.get_ext_fullpath(ext.name))
        ext_dir.mkdir(parents=True, exist_ok=True)

        config = "Debug" if self.debug else "Release"

        all_args = [
            # "-DCMAKE_LIBRARY_OUTPUT_DIRECTORY=" + str(ext_dir.parent.absolute()),
            "-DCMAKE_INSTALL_PREFIX=" + str(ext_dir.parent.absolute()),
            "-DCMAKE_BUILD_TYPE=" + config,
            "-DPYBIND_LIB=" + "ON",
        ]
        if len(self.cmake_args.strip()) > 0:
            for argv in self.cmake_args.split():
                all_args.append(argv)

        build_args = [
            "--", "-j32"
        ]

        os.chdir(build_temp)
        cmake_cmd = ["cmake", f"{str(cwd)}/{ext.name}"] + all_args
        self.spawn(cmake_cmd)
        if not self.dry_run:
            cmake_cmd = ["cmake", "--build", "."] + build_args
            self.spawn(cmake_cmd)
        if not self.dry_run:
            cmake_cmd = ["cmake", "--install", "."]
            self.spawn(cmake_cmd)
        os.chdir(str(cwd))


def get_requirements() -> List[str]:
    """Get Python package dependencies from requirements.txt."""

    def _read_requirements(filename: str) -> List[str]:
        with open(os.path.join(ROOT_DIR, filename)) as f:
            requirements = f.read().strip().split("\n")
        resolved_requirements = []
        for line in requirements:
            if line.startswith("-r "):
                resolved_requirements += _read_requirements(line.split()[1])
            else:
                resolved_requirements.append(line)
        return resolved_requirements

    if _is_npu():
        requirements = _read_requirements("requirements-cpu.txt")
    else:
        raise ValueError(
            "Unsupported platform, please use CUDA, ROCm, Neuron, or CPU.")
    return requirements


setup(
    name="txda",
    version="0.0.3",
    author="zhoushen",
    author_email="",
    license="Apache 2.0",
    description="A Python Library for AI Accelerator Programming Language and Runtime",
    long_description="",
    classifiers=[
        "Programming Language :: Python :: 3.8",
        "Programming Language :: Python :: 3.9",
        "Programming Language :: Python :: 3.10",
        "Programming Language :: Python :: 3.11",
        "License :: OSI Approved :: Apache Software License",
        "Topic :: Scientific/Engineering :: Artificial Intelligence",
    ],
    packages=find_packages(exclude=("examples", "csrc", "cmake", "third-party/tx8-host", "third-party/tx8-yoc-rt-thread-smp")),
    install_requires=[], # get_requirements(),
    ext_modules=[CMakeExtension(".")],
    # package_data={
    #     '': ['lib/*.*'],
    # },
    cmdclass=dict(build_ext=CMakeBuild),
    zip_safe=False,
    python_requires=">=3.8",
)
